import pymysql
import os

class AuthManager:
    def __init__(self):
        self.mysql_pwd = os.getenv('MYSQL_PWD', "123456")
        self.db_name = os.getenv('DB_NAME', "appDB")

    def get_db_connection(self):
        """获取数据库连接"""
        return pymysql.connect(
            host="localhost",
            user="root",
            password=self.mysql_pwd,
            database=self.db_name,
            charset='utf8'
        )

    def register_user(self, username, password, phone, addr, user_role):
        """用户注册"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()

            # 根据角色选择表
            if user_role == 'RESTAURANT':
                table = 'RESTAURANT'
            elif user_role == 'CUSTOMER':
                table = 'CUSTOMER'
            else:
                return "无效的用户角色。"

            # 检查用户是否已存在
            sql_check = f"SELECT * FROM {table} WHERE username = %s"
            cursor.execute(sql_check, (username,))
            if cursor.fetchone():
                return "用户已注册，请直接登录。"

            # 插入新用户
            sql_insert = f"INSERT INTO {table} (username, password, address, phone) VALUES (%s, %s, %s, %s)"
            cursor.execute(sql_insert, (username, password, addr, phone))
            db.commit()
            return "注册成功！"
        except Exception as e:
            print(f"Error: {e}")
            return "注册失败，请稍后重试。"
        finally:
            if db:
                db.close()

    def login_user(self, username, password, user_role):
        """用户登录"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()

            # 根据角色选择表
            if user_role == 'ADMIN':
                table = 'ADMIN'
            elif user_role == 'RESTAURANT':
                table = 'RESTAURANT'
            elif user_role == 'CUSTOMER':
                table = 'CUSTOMER'
            else:
                return "无效的用户角色。"

            # 验证用户名和密码
            sql_check = f"SELECT * FROM {table} WHERE username = %s AND password = %s"
            cursor.execute(sql_check, (username, password))
            if cursor.fetchone():
                return f"登录成功！欢迎 {user_role} 用户！"
            else:
                return "用户名或密码错误。"
        except Exception as e:
            print(f"Error: {e}")
            return "登录失败，请稍后重试。"
        finally:
            if db:
                db.close()